import os
# -*- coding: utf-8 -*-
import einops
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import transforms
from pathlib import Path
from PIL import Image


class MyDataset(Dataset):
    def __init__(self, root):
        '''

        :param root: image path
        '''
        self.transform = transforms.Compose([
            transforms.Resize([128,128]),
            # transforms.CenterCrop(128),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        images_path = Path(root)
        images_list = list(images_path.glob('*.jpg'))
        images_list_str = [str(x) for x in images_list]
        self.images = images_list_str


    def __getitem__(self, item):
        image_path = self.images[item]
        image = Image.open(image_path)
        return self.transform(image)

    def __len__(self):
        return len(self.images)


mydataloader = DataLoader(
    MyDataset(r'C:\Users\86186\Desktop\study\大三下\智能计算系统\PROJECT\DATA\Market-1501-v15.09.15\bounding_box_train'),
    batch_size=64, shuffle=False)

if __name__ == '__main__':
    mydataloader = DataLoader(
        MyDataset(r'C:\Users\86186\Desktop\study\大三下\智能计算系统\PROJECT\DATA\Market-1501-v15.09.15\bounding_box_train'),
        batch_size=1, shuffle=False)
    for i_batch, batch_data in enumerate(mydataloader):
        print(i_batch)
        print(batch_data.shape)
        show_img = einops.rearrange(batch_data, 'b c h w->h (b w) c').numpy()
        show_img = np.uint8(show_img*255)
        from PIL import Image

        print(show_img.shape)
        pilimg = Image.fromarray(show_img)
        pilimg.show()

        break
